#!/bin/bash

####################################################################
# 查看内存带宽，适用于Kunpeng 920/920x
#
# 使用方式：
#     ./mem_bw
#
# 更新说明
# - 2025.08: 支持920x，支持GB/KB单位，默认持续运行
# - 2025.02: 支持interval > 1的情况；使用1000而不是1024计算MB
####################################################################

######################### Global Variables #########################

DURATION=0          # sampling duration in seconds
INTERVAL=1          # reporting interval in seconds
UNIT="MB"           # unit for output bandwidth

####################################################################

sample_bandwidth_920() {
    # 内存带宽事件
    local cpu0_die0_rd="hisi_sccl1_ddrc0/flux_rd/,hisi_sccl1_ddrc1/flux_rd/,hisi_sccl1_ddrc2/flux_rd/,hisi_sccl1_ddrc3/flux_rd/"
    local cpu0_die0_wr="hisi_sccl1_ddrc0/flux_wr/,hisi_sccl1_ddrc1/flux_wr/,hisi_sccl1_ddrc2/flux_wr/,hisi_sccl1_ddrc3/flux_wr/"
    local cpu0_die1_rd="hisi_sccl3_ddrc0/flux_rd/,hisi_sccl3_ddrc1/flux_rd/,hisi_sccl3_ddrc2/flux_rd/,hisi_sccl3_ddrc3/flux_rd/"
    local cpu0_die1_wr="hisi_sccl3_ddrc0/flux_wr/,hisi_sccl3_ddrc1/flux_wr/,hisi_sccl3_ddrc2/flux_wr/,hisi_sccl3_ddrc3/flux_wr/"
    local cpu1_die0_rd="hisi_sccl5_ddrc0/flux_rd/,hisi_sccl5_ddrc1/flux_rd/,hisi_sccl5_ddrc2/flux_rd/,hisi_sccl5_ddrc3/flux_rd/"
    local cpu1_die0_wr="hisi_sccl5_ddrc0/flux_wr/,hisi_sccl5_ddrc1/flux_wr/,hisi_sccl5_ddrc2/flux_wr/,hisi_sccl5_ddrc3/flux_wr/"
    local cpu1_die1_rd="hisi_sccl7_ddrc0/flux_rd/,hisi_sccl7_ddrc1/flux_rd/,hisi_sccl7_ddrc2/flux_rd/,hisi_sccl7_ddrc3/flux_rd/"
    local cpu1_die1_wr="hisi_sccl7_ddrc0/flux_wr/,hisi_sccl7_ddrc1/flux_wr/,hisi_sccl7_ddrc2/flux_wr/,hisi_sccl7_ddrc3/flux_wr/"

    local ddrc_events=" $cpu0_die0_rd,$cpu0_die0_wr,$cpu0_die1_rd,$cpu0_die1_wr,$cpu1_die0_rd,$cpu1_die0_wr,$cpu1_die1_rd,$cpu1_die1_wr "

    # 检查内存带宽事件是否全部合法
    if ! perf stat -e ${ddrc_events} --timeout 10 &>/dev/null ; then
        echo -e "Unsupported PMU event. See below for details:\n"
        perf stat -e ${ddrc_events} --timeout 10
        exit 1
    fi

    local timeout_prefix=""
    if [[ $DURATION -ne 0 ]]; then
        timeout_prefix="timeout $DURATION"
    fi

    local interval=$(($INTERVAL * 1000))

    $timeout_prefix perf stat -x ',' -a -I $(($INTERVAL * 1000)) -e ${ddrc_events} 2>&1 | \
	awk -v interval=$INTERVAL -v unit=$UNIT -F, '
		BEGIN {
			tpre = -1;
			first = 1;

            if (unit == "GB") {
                factor = 32 / 1000 / 1000 / 1000 / interval;
            } else if (unit == "KB") {
                factor = 32 / 1000 / interval;
            } else {
                factor = 32 / 1000 / 1000 / interval;
            }
		}
		{
			tcurr = $1;
			if (first == 1) {
				tpre = tcurr;
				first = 0;
			}

			if (tcurr != tpre) {
                printf("\nSystem DDR Bandwidth at %s\n\n", strftime("%Y-%m-%d %H:%M:%S"));

                printf("                                CPU0             |              CPU1             |    TOTAL\n");
                printf("                       NUMA0     |     NUMA1     |     NUMA2     |     NUMA3     |\n");
                printf("DDR RD    BW :     %9.2f     | %9.2f     | %9.2f     | %9.2f     | %8.2f %s/s\n",
                    CPU0_Die0_R * factor,
                    CPU0_Die1_R * factor,
                    CPU1_Die0_R * factor,
                    CPU1_Die1_R * factor,
                    (CPU0_Die0_R + CPU0_Die1_R + CPU1_Die0_R + CPU1_Die1_R) * factor,
                    unit);
                printf("DDR WR    BW :     %9.2f     | %9.2f     | %9.2f     | %9.2f     | %8.2f %s/s\n",
                    CPU0_Die0_W * factor,
                    CPU0_Die1_W * factor,
                    CPU1_Die0_W * factor,
                    CPU1_Die1_W * factor,
                    (CPU0_Die0_W + CPU0_Die1_W + CPU1_Die0_W + CPU1_Die1_W) * factor,
                    unit);
                printf("DDR TOTAL BW :     %9.2f     | %9.2f     | %9.2f     | %9.2f     | %8.2f %s/s\n",
                    (CPU0_Die0_R + CPU0_Die0_W) * factor,
                    (CPU0_Die1_R + CPU0_Die1_W) * factor,
                    (CPU1_Die0_R + CPU1_Die0_W) * factor,
                    (CPU1_Die1_R + CPU1_Die1_W) * factor,
                    (CPU0_Die0_R + CPU0_Die1_R + CPU1_Die0_R + CPU1_Die1_R + CPU0_Die0_W + CPU0_Die1_W + CPU1_Die0_W + CPU1_Die1_W) * factor,
                    unit);

				CPU0_Die0_R=0;
				CPU0_Die0_W=0;
				CPU0_Die1_R=0;
				CPU0_Die1_W=0;
				CPU1_Die0_R=0;
				CPU1_Die0_W=0;
				CPU1_Die1_R=0;
				CPU1_Die1_W=0;

				tpre = tcurr;
			}

			switch ($4) {
				case /sccl1.*rd/:
					CPU0_Die0_R += $2;
					break;
				case /sccl1.*wr/:
					CPU0_Die0_W += $2;
					break;
				case /sccl3.*rd/:
					CPU0_Die1_R += $2;
					break;
				case /sccl3.*wr/:
					CPU0_Die1_W += $2;
					break;
				case /sccl5.*rd/:
					CPU1_Die0_R += $2;
					break;
				case /sccl5.*wr/:
					CPU1_Die0_W += $2;
					break;
				case /sccl7.*rd/:
					CPU1_Die1_R += $2;
					break;
				case /sccl7.*wr/:
					CPU1_Die1_W += $2;
					break;
			}
		}
	'
}

sample_bandwidth_920x() {
    # 内存带宽事件
    local cpu0_die0_rd="hisi_sccl3_ddrc0_0/flux_rd/,hisi_sccl3_ddrc0_1/flux_rd/,hisi_sccl3_ddrc2_0/flux_rd/,hisi_sccl3_ddrc2_1/flux_rd/,hisi_sccl3_ddrc3_0/flux_rd/,hisi_sccl3_ddrc3_1/flux_rd/,hisi_sccl3_ddrc5_0/flux_rd/,hisi_sccl3_ddrc5_1/flux_rd/"
    local cpu0_die0_wr="hisi_sccl3_ddrc0_0/flux_wr/,hisi_sccl3_ddrc0_1/flux_wr/,hisi_sccl3_ddrc2_0/flux_wr/,hisi_sccl3_ddrc2_1/flux_wr/,hisi_sccl3_ddrc3_0/flux_wr/,hisi_sccl3_ddrc3_1/flux_wr/,hisi_sccl3_ddrc5_0/flux_wr/,hisi_sccl3_ddrc5_1/flux_wr/"
    local cpu0_die1_rd="hisi_sccl1_ddrc0_0/flux_rd/,hisi_sccl1_ddrc0_1/flux_rd/,hisi_sccl1_ddrc2_0/flux_rd/,hisi_sccl1_ddrc2_1/flux_rd/,hisi_sccl1_ddrc3_0/flux_rd/,hisi_sccl1_ddrc3_1/flux_rd/,hisi_sccl1_ddrc5_0/flux_rd/,hisi_sccl1_ddrc5_1/flux_rd/"
    local cpu0_die1_wr="hisi_sccl1_ddrc0_0/flux_wr/,hisi_sccl1_ddrc0_1/flux_wr/,hisi_sccl1_ddrc2_0/flux_wr/,hisi_sccl1_ddrc2_1/flux_wr/,hisi_sccl1_ddrc3_0/flux_wr/,hisi_sccl1_ddrc3_1/flux_wr/,hisi_sccl1_ddrc5_0/flux_wr/,hisi_sccl1_ddrc5_1/flux_wr/"
    local cpu1_die0_rd="hisi_sccl11_ddrc0_0/flux_rd/,hisi_sccl11_ddrc0_1/flux_rd/,hisi_sccl11_ddrc2_0/flux_rd/,hisi_sccl11_ddrc2_1/flux_rd/,hisi_sccl11_ddrc3_0/flux_rd/,hisi_sccl11_ddrc3_1/flux_rd/,hisi_sccl11_ddrc5_0/flux_rd/,hisi_sccl11_ddrc5_1/flux_rd/"
    local cpu1_die0_wr="hisi_sccl11_ddrc0_0/flux_wr/,hisi_sccl11_ddrc0_1/flux_wr/,hisi_sccl11_ddrc2_0/flux_wr/,hisi_sccl11_ddrc2_1/flux_wr/,hisi_sccl11_ddrc3_0/flux_wr/,hisi_sccl11_ddrc3_1/flux_wr/,hisi_sccl11_ddrc5_0/flux_wr/,hisi_sccl11_ddrc5_1/flux_wr/"
    local cpu1_die1_rd="hisi_sccl9_ddrc0_0/flux_rd/,hisi_sccl9_ddrc0_1/flux_rd/,hisi_sccl9_ddrc2_0/flux_rd/,hisi_sccl9_ddrc2_1/flux_rd/,hisi_sccl9_ddrc3_0/flux_rd/,hisi_sccl9_ddrc3_1/flux_rd/,hisi_sccl9_ddrc5_0/flux_rd/,hisi_sccl9_ddrc5_1/flux_rd/"
    local cpu1_die1_wr="hisi_sccl9_ddrc0_0/flux_wr/,hisi_sccl9_ddrc0_1/flux_wr/,hisi_sccl9_ddrc2_0/flux_wr/,hisi_sccl9_ddrc2_1/flux_wr/,hisi_sccl9_ddrc3_0/flux_wr/,hisi_sccl9_ddrc3_1/flux_wr/,hisi_sccl9_ddrc5_0/flux_wr/,hisi_sccl9_ddrc5_1/flux_wr/"

    local ddrc_events=" $cpu0_die0_rd,$cpu0_die0_wr,$cpu0_die1_rd,$cpu0_die1_wr,$cpu1_die0_rd,$cpu1_die0_wr,$cpu1_die1_rd,$cpu1_die1_wr "

    # 检查内存带宽事件是否全部合法
    if ! perf stat -e ${ddrc_events} --timeout 10 &>/dev/null ; then
        echo -e "Unsupported PMU event. See below for details:\n"
        perf stat -e ${ddrc_events} --timeout 10
        exit 1
    fi

    local timeout_prefix=""
    if [[ $DURATION -ne 0 ]]; then
        timeout_prefix="timeout $DURATION"
    fi

    $timeout_prefix perf stat -x ',' -a -I $(($INTERVAL * 1000)) -e ${ddrc_events} 2>&1 | \
	awk -v interval=$INTERVAL -v unit=$UNIT -F, '
		BEGIN {
			tpre = -1;
			first = 1;

            if (unit == "GB") {
                factor = 32 / 1000 / 1000 / 1000 / interval;
            } else if (unit == "KB") {
                factor = 32 / 1000 / interval;
            } else {
                factor = 32 / 1000 / 1000 / interval;
            }
		}
		{
			tcurr = $1;
			if (first == 1) {
				tpre = tcurr;
				first = 0;
			}

			if (tcurr != tpre) {
                printf("\nSystem DDR Bandwidth at %s\n\n", strftime("%Y-%m-%d %H:%M:%S"));

                printf("                                CPU0             |              CPU1             |    TOTAL\n");
                printf("                       NUMA0     |     NUMA1     |     NUMA2     |     NUMA3     |\n");
                printf("DDR RD    BW :     %9.2f     | %9.2f     | %9.2f     | %9.2f     | %8.2f %s/s\n",
                    CPU0_Die0_R * factor,
                    CPU0_Die1_R * factor,
                    CPU1_Die0_R * factor,
                    CPU1_Die1_R * factor,
                    (CPU0_Die0_R + CPU0_Die1_R + CPU1_Die0_R + CPU1_Die1_R) * factor,
                    unit);
                printf("DDR WR    BW :     %9.2f     | %9.2f     | %9.2f     | %9.2f     | %8.2f %s/s\n",
                    CPU0_Die0_W * factor,
                    CPU0_Die1_W * factor,
                    CPU1_Die0_W * factor,
                    CPU1_Die1_W * factor,
                    (CPU0_Die0_W + CPU0_Die1_W + CPU1_Die0_W + CPU1_Die1_W) * factor,
                    unit);
                printf("DDR TOTAL BW :     %9.2f     | %9.2f     | %9.2f     | %9.2f     | %8.2f %s/s\n",
                    (CPU0_Die0_R + CPU0_Die0_W) * factor,
                    (CPU0_Die1_R + CPU0_Die1_W) * factor,
                    (CPU1_Die0_R + CPU1_Die0_W) * factor,
                    (CPU1_Die1_R + CPU1_Die1_W) * factor,
                    (CPU0_Die0_R + CPU0_Die1_R + CPU1_Die0_R + CPU1_Die1_R + CPU0_Die0_W + CPU0_Die1_W + CPU1_Die0_W + CPU1_Die1_W) * factor,
                    unit);

				CPU0_Die0_R=0;
				CPU0_Die0_W=0;
				CPU0_Die1_R=0;
				CPU0_Die1_W=0;
				CPU1_Die0_R=0;
				CPU1_Die0_W=0;
				CPU1_Die1_R=0;
				CPU1_Die1_W=0;

				tpre = tcurr;
			}

			switch ($4) {
                case /sccl3_.*rd/:
					CPU0_Die0_R += $2;
					break;
				case /sccl3_.*wr/:
					CPU0_Die0_W += $2;
					break;
				case /sccl1_.*rd/:
					CPU0_Die1_R += $2;
					break;
				case /sccl1_.*wr/:
					CPU0_Die1_W += $2;
					break;
				case /sccl11_.*rd/:
					CPU1_Die0_R += $2;
					break;
				case /sccl11_.*wr/:
					CPU1_Die0_W += $2;
					break;
				case /sccl9_.*rd/:
					CPU1_Die1_R += $2;
					break;
				case /sccl9_.*wr/:
					CPU1_Die1_W += $2;
					break;
			}
		}
	'
}

help() {
    cat <<- EOF
Display DDRC bandwidth

Usage:
    $0

Options:
    -d, --duration      Sampling duration in seconds. Default running continuously.
    -i, --interval      Reporting interval in seconds. Default: 1

    -g, --gb            Show output in GB (10^9). If not specified, MB is used.
    -m, --mb            Show output in MB (10^6). Default.
    -k, --kb            Show output in KB (10^3). If not specified, MB is used.

    -h, --help          Show usage.

Example:
    每采样5秒打印一次，共采样60秒
    $0 -i 5 -d 60

    按GB显示
    $0 -g

EOF
}

main() {
    # 检查GNU getopt兼容性
    getopt -T >/dev/null 2>&1
    local ret=$?
    if [ $ret -ne 4 ]; then
        echo "Error: This script requires GNU getopt for long option support." >&2
        exit 1
    fi

    # 定义短选项和长选项
    local short_opts="d:i:gmkh"
    local long_opts="duration:,interval:,gb,mb,kb,help"

    # 使用getopt解析参数
    local parsed
    parsed=$(getopt -o "$short_opts" -l "$long_opts" -n "$0" -- "$@")

    if [ $? -ne 0 ]; then
        help
        exit 1
    fi
    eval set -- "$parsed"

    # 处理解析后的选项
    while true; do
        case "$1" in
            -d | --duration )
                DURATION="$2"
                shift 2
                ;;
            -i | --interval )
                INTERVAL="$2"
                shift 2
                ;;
            -g | --gb )
                UNIT="GB"
                shift 1
                ;;
            -m | --mb )
                UNIT="MB"
                shift 1
                ;;
            -k | --kb )
                UNIT="KB"
                shift 1
                ;;
            -h | --help )
                help
                exit 0
                ;;
            --)
                shift
                break
                ;;
            *)
                echo "Error: Internal error!" >&2
                exit 1
                ;;
        esac
    done

    # 检查是否Kunpeng处理器
    if ! dmidecode -t processor | grep -q 'Kunpeng' ; then
		echo "This CPU Architecture is not supported. Please contact the author."
		exit 1
    fi

    if perf stat -e hisi_sccl9_ddrc5_0/flux_rd/ --timeout 10 &>/dev/null ; then
        sample_bandwidth_920x
    else
        sample_bandwidth_920
    fi
}

main "$@"
